import time
from typing import Any

import cv2
import numpy
from insightface.app import FaceAnalysis
from insightface.model_zoo.inswapper import INSwapper

from Settings import ModelsBaseDIR


class ChangeFace:
    def __init__(self):
        self.__FaceAnalysis = self.__initFaceAnalysisEngine()
        self.__FaceSwapper = self.__initFaceSwapperEngine()

    def __initFaceAnalysisEngine(self):
        face_analysis = FaceAnalysis(
            # 下载模型的目录名称
            name='buffalo_l',
            # 模型下载位置
            root=ModelsBaseDIR,
            # 'AzureExecutionProvider, CPUExecutionProvider'
            # providers=["CPUExecutionProvider"]
        )
        # 预处理
        face_analysis.prepare(ctx_id=0, det_size=(640, 640))

        return face_analysis

    def __initFaceSwapperEngine(self):
        inswapper = INSwapper(
            model_file=fr'C:\Users\Administrator\Desktop\change_face\change_face\models\1.onnx',
        )

        return inswapper

    def _getOneFace(self, img_path) -> numpy.ndarray:
        if isinstance(img_path, str):
            t_frame = cv2.imread(img_path)
        else:
            t_frame = img_path
        face = self.__FaceAnalysis.get(t_frame)
        try:
            # 如果拿到了多张脸，就只返回最近的一个，返回了一个Face的对象
            return min(face, key=lambda x: x.bbox[0])
        except ValueError:
            raise "传入图片地址"


    def _swapFace(self, tmp_face, target_face, source_face) -> numpy.ndarray[Any, Any]:
        frame = self.__FaceSwapper.get(tmp_face, target_face, source_face, paste_back=True)
        return frame

    def start(self, target_path, source_path):
        source_face = self._getOneFace(source_path)
        target_face = self._getOneFace(target_path)

        # tmp_region_result: numpy.ndarray[Any, Any] = cv2.imread(target_path)
        tmp_region_result: numpy.ndarray[Any, Any] = target_path

        frame = self._swapFace(tmp_region_result, target_face, source_face)

        return frame

if __name__ == '__main__':
    from Settings import BaseTestDataCubePath, OutBaseDir

    changeFace = ChangeFace()
    camera = cv2.VideoCapture(0)  # Use index for the webcam (adjust the index accordingly if necessary)
    camera.set(cv2.CAP_PROP_FRAME_WIDTH, 960)  # Set the width of the resolution
    camera.set(cv2.CAP_PROP_FRAME_HEIGHT, 540)  # Set the height of the resolution
    camera.set(cv2.CAP_PROP_FPS, 60)

    target_image_path = f'{BaseTestDataCubePath}\\4.jpg'



    while camera:
        ret, frame = camera.read()

        change_frame = changeFace.start(frame, target_image_path)
        if ret:
            cv2.imshow('Camera', change_frame)




        # 按下 'q' 键退出循环
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    # 释放摄像头资源并关闭窗口
    camera.release()
    cv2.destroyAllWindows()